#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GitCode 抓包脚本
功能：自动登录 -> 获取 token -> 记录所有请求 -> 保存到文件
"""

import json
import time
import threading
import os
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager

class GitCodeCapture:
    def __init__(self):
        self.username = "fengya011"
        self.password = "Ge.19920204"
        self.login_url = "https://gitcode.com/dashboard"
        self.all_requests = []
        self.cookies = []
        self.driver = None
        self.is_capturing = False
        
    def setup_driver(self):
        """设置 Chrome 浏览器"""
        print("正在启动 Chrome 浏览器...")
        
        options = Options()
        options.add_argument("--disable-blink-features=AutomationControlled")
        options.add_argument("--disable-web-security")
        options.add_argument("--disable-features=VizDisplayCompositor")
        options.add_experimental_option("excludeSwitches", ["enable-automation"])
        options.add_experimental_option('useAutomationExtension', False)
        
        # 启用性能日志
        options.set_capability("goog:loggingPrefs", {"performance": "ALL"})
        
        try:
            # 自动下载和管理 chromedriver
            self.driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
            self.driver.execute_script("Object.defineProperty(navigator, 'webdriver', {get: () => undefined})")
            
            # 启用网络监听
            self.driver.execute_cdp_cmd('Network.enable', {})
            self.driver.execute_cdp_cmd('Performance.enable', {})
            
            print("✅ Chrome 浏览器启动成功")
            return True
        except Exception as e:
            print(f"❌ Chrome 浏览器启动失败: {e}")
            print("请确保已安装 Chrome 浏览器")
            return False
    
    def login(self):
        """登录 GitCode"""
        print("正在登录 GitCode...")
        
        try:

             # 2. 自动跳转到登录页
            self.driver.get(self.login_url)
            time.sleep(3)
            # 等待所有登录tab加载
            wait = WebDriverWait(self.driver, 10)
            tabs = wait.until(
                EC.presence_of_all_elements_located((By.CLASS_NAME, "login-modal-tab"))
            )
            password_tab = None
            for tab in tabs:
                if "密码登录" in tab.text:
                    password_tab = tab
                    # 如果不是激活状态，点击它
                    if "login-modal-tab-active" not in tab.get_attribute("class"):
                        tab.click()
                        print("已切换到密码登录")
                        time.sleep(1)
                    else:
                        print("已在密码登录页签")
                    break
            if not password_tab:
                print("未找到密码登录标签")
            
            # 勾选同意协议复选框（如果未勾选）
            try:
                agree_checkbox = self.driver.find_element(By.CSS_SELECTOR, "input[type='checkbox']")
                if not agree_checkbox.is_selected():
                    agree_checkbox.click()
                    print("已勾选同意协议复选框")
            except Exception as e:
                print(f"未找到或无需勾选同意协议复选框: {e}")

            # 查找并填写用户名
            username_input = wait.until(
                EC.presence_of_element_located((By.XPATH, "//input[@placeholder='请填写用户名/邮箱']"))
            )
            username_input.clear()
            username_input.send_keys(self.username)
            print(f"已输入用户名: {self.username}")
            
            # 查找并填写密码
            password_input = wait.until(
                EC.presence_of_element_located((By.XPATH, "//input[@placeholder='请填写密码']"))
            )
            password_input.clear()
            password_input.send_keys(self.password)
            print("已输入密码")
            
            # 点击登录按钮或按回车
            password_input.send_keys(Keys.RETURN)
            print("已提交登录表单")
            
            # 等待登录完成
            time.sleep(5)
            
            # 检查登录是否成功
            current_url = self.driver.current_url
            if "dashboard" in current_url or "profile" in current_url:
                print("✅ 登录成功！")
                return True
            else:
                print(f"⚠️ 登录状态不明确，当前页面: {current_url}")
                # 检查是否有错误信息
                page_source = self.driver.page_source
                if "用户名或密码错误" in page_source or "Invalid" in page_source:
                    print("❌ 登录失败：用户名或密码错误")
                    return False
                return True
                
        except Exception as e:
            print(f"❌ 登录过程中出现错误: {e}")
            return False
    
    def start_capture(self):
        """开始抓取请求"""
        print("开始抓取网络请求...")
        self.is_capturing = True
        
        def capture_requests():
            # 清空历史日志，只抓取新请求
            self.driver.get_log("performance")
            
            while self.is_capturing:
                try:
                    # 获取性能日志
                    logs = self.driver.get_log("performance")
                    for entry in logs:
                        # 原始日志写入
                        with open("raw_cdp_log.txt", "a", encoding="utf-8") as logf:
                            logf.write(str(entry) + "\n")
                        # 结构化解析
                        try:
                            message = json.loads(entry["message"])['message']
                            if message["method"] == "Network.requestWillBeSent":
                                req = message["params"]["request"]
                                req_info = {
                                    "url": req.get("url"),
                                    "method": req.get("method"),
                                    "headers": req.get("headers"),
                                    "postData": req.get("postData", ""),
                                    "timestamp": message["params"].get("timestamp")
                                }
                                self.all_requests.append(req_info)
                        except Exception:
                            pass
                    time.sleep(0.1)  # 短暂休眠，避免过度占用CPU
                except Exception as e:
                    print(f"抓取请求时出现错误: {e}")
                    break
        
        # 启动抓取线程
        self.capture_thread = threading.Thread(target=capture_requests)
        self.capture_thread.daemon = True
        self.capture_thread.start()
    
    def stop_capture(self):
        """停止抓取"""
        print("正在停止抓取...")
        self.is_capturing = False
        if hasattr(self, 'capture_thread'):
            self.capture_thread.join(timeout=2)
    
    def get_cookies_and_tokens(self):
        """获取 cookies 和 tokens"""
        print("正在获取 cookies 和 tokens...")
        
        try:
            # 获取所有 cookies
            self.cookies = self.driver.get_cookies()
            
            # 获取 localStorage
            local_storage = self.driver.execute_script("""
                var items = {};
                for (var i = 0; i < localStorage.length; i++) {
                    var key = localStorage.key(i);
                    items[key] = localStorage.getItem(key);
                }
                return items;
            """)
            
            # 获取 sessionStorage
            session_storage = self.driver.execute_script("""
                var items = {};
                for (var i = 0; i < sessionStorage.length; i++) {
                    var key = sessionStorage.key(i);
                    items[key] = sessionStorage.getItem(key);
                }
                return items;
            """)
            
            # 尝试获取 CSRF token
            csrf_token = self.driver.execute_script("""
                var token = document.querySelector('meta[name="csrf-token"]');
                return token ? token.getAttribute('content') : null;
            """)
            
            return {
                "cookies": self.cookies,
                "localStorage": local_storage,
                "sessionStorage": session_storage,
                "csrf_token": csrf_token
            }
            
        except Exception as e:
            print(f"获取 cookies 和 tokens 时出现错误: {e}")
            return {}
    
    def save_to_file(self, filename="gitcode_capture_data.json"):
        """保存数据到文件"""
        print(f"正在保存数据到 {filename}...")
        
        try:
            # 获取 tokens
            tokens_data = self.get_cookies_and_tokens()
            
            # 整理数据
            capture_data = {
                "capture_time": datetime.now().isoformat(),
                "username": self.username,
                "login_url": self.login_url,
                "tokens": tokens_data,
                "requests": self.all_requests,
                "total_requests": len(self.all_requests)
            }
            
            # 保存到文件（覆盖模式）
            with open(filename, "w", encoding="utf-8") as f:
                json.dump(capture_data, f, ensure_ascii=False, indent=2)
            
            print(f"✅ 数据已保存到 {filename}")
            print(f"📊 统计信息:")
            print(f"   - 总请求数: {len(self.all_requests)}")
            print(f"   - Cookies 数量: {len(self.cookies)}")
            print(f"   - 抓取时间: {capture_data['capture_time']}")
            
            return True
            
        except Exception as e:
            print(f"❌ 保存文件时出现错误: {e}")
            return False
    
    def run(self):
        """运行完整的抓包流程（支持交互式开始/暂停/终止）"""
        print("🚀 GitCode 抓包脚本启动")
        print("=" * 50)
        try:
            # 1. 启动浏览器
            if not self.setup_driver():
                return False
            
            self.login()
           
            print("💡 你可以在浏览器中手动操作页面。")
            print("指令：s=开始抓包，p=暂停抓包，q=终止并保存")
            capturing = False
            while True:
                cmd = input("请输入指令 (s=开始, p=暂停, q=退出)：").strip().lower()
                if cmd == "s":
                    if not capturing:
                        print("🔍 开始抓包...")
                        self.start_capture()
                        capturing = True
                    else:
                        print("⚠️ 已经在抓包中，无需重复开始。")
                elif cmd == "p":
                    if capturing:
                        print("⏸️ 暂停抓包...")
                        self.stop_capture()
                        capturing = False
                    else:
                        print("⚠️ 当前未在抓包，无需暂停。")
                elif cmd == "q":
                    if capturing:
                        print("⏹️ 停止抓包...")
                        self.stop_capture()
                        capturing = False
                    print("💾 正在保存数据...")
                    self.save_to_file()
                    self.driver.quit()
                    print("✅ 抓包流程结束")
                    break
                else:
                    print("无效指令，请输入 s、p 或 q。")
            return True
        except Exception as e:
            print(f"❌ 脚本运行异常: {e}")
            return False

def main():
    """主函数"""
    capture = GitCodeCapture()
    capture.run()

if __name__ == "__main__":
    main() 